#!/usr/bin/python

import socket
import struct
import time
import binascii
import json
import numpy as np
import scipy.interpolate as sci
import cv2
import matplotlib.pyplot as plt
import re


numTxAzimAnt = 2
numRxAnt = 4
numRangeBins = 256
NUM_ANGLE_BINS = 64
freqSlopeConst_actual = 70 #会变化要注意
digOutSampleRate = 5209 #会变化要注意
rangeBias = 0  #会变化要注意
range_width = 5
range_depth = 10
rangeAzimuthHeatMapGrid_points = 100

MaxBytes=4096

server = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
server.settimeout(600)

host_server_GetData = ''
port_server_GetData = 5678
server.bind((host_server_GetData, port_server_GetData)) 



def main():
    #start
    #start = time.clock()
    while(1):
        
        # time.sleep(2)

        start = time.perf_counter()

        numBytes = numTxAzimAnt * numRxAnt * numRangeBins * 4

        qrows = numTxAzimAnt * numRxAnt
        qcols = numRangeBins
        qidx = 0
        QQ_ = np.zeros((256,64),dtype=float)
        data = []
        raw = np.zeros(NUM_ANGLE_BINS, dtype=complex)
        raw_ = np.zeros(NUM_ANGLE_BINS, dtype=complex)
        data_raw = []

        INPUT_data = ""
        INPUT_data = tcpserver_Getdata()

        if (len(INPUT_data) < 16300):
            continue


        if(INPUT_data.find("0400000000200000") != -1):
            data_raw = INPUT_data[INPUT_data.find("0400000000200000")+len("0400000000200000"):INPUT_data.find("0400000000200000")+len("0400000000200000")+2*numBytes]
            data = re.findall(r'\w{2}', data_raw)
            length = len(data)
            for i in range(length):
                data[i] = int(data[i],16)
        else:
            print("Wrong")

        ### 这里数制转化存在问题
        # 2 拆分字符串数组？

        for tmpc in range(qcols):
            real = np.zeros(NUM_ANGLE_BINS)
            imag = np.zeros(NUM_ANGLE_BINS)
            for tmpr in range(qrows):
                real[tmpr] = data[qidx + 1] * 256 + data[qidx]
                if (real[tmpr] > 32767): real[tmpr] = real[tmpr] - 65536
                imag[tmpr] = data[qidx + 3] * 256 + data[qidx + 2]
                if (imag[tmpr] > 32767): imag[tmpr] = imag[tmpr] - 65536
                qidx = qidx + 4

            for i in range(NUM_ANGLE_BINS):
                raw[i] = complex(real[i], -imag[i])

            # 出品的一维与fft2一致
            raw_ = np.fft.fft(raw)
            raw_real = np.abs(raw_)
            # QQ
            a = raw_real[32:]
            b = raw_real[0:32]
            QQ_[tmpc] = np.append(a,b)

        range_1 = np.arange(-NUM_ANGLE_BINS / 2 + 1, NUM_ANGLE_BINS / 2)

        theta = np.arcsin(np.dot(range_1, 2 / NUM_ANGLE_BINS))

        rangeIdxToMeters = 3e8 * digOutSampleRate * 1e3 / (2 * np.abs(freqSlopeConst_actual) * 1e12 * numRangeBins)
        range_2 = np.arange(0, numRangeBins)
        # rangeIdxToMeters可能需要根据使用场景不同进行变换
        range_d = np.dot(range_2, rangeIdxToMeters)
        range_d = np.subtract(range_d, rangeBias)

        for i in range(numRangeBins):
            range_d[i] = np.max(range_d[i], 0)

        posX_ = np.outer(range_d, np.sin(theta))
        posY_ = np.outer(range_d, np.cos(theta))

        rangeAzimuthHeatMapGrid_xlin = np.arange(-range_width, range_width, 2.0 * range_width / (rangeAzimuthHeatMapGrid_points - 1))
        if (len(rangeAzimuthHeatMapGrid_xlin) < rangeAzimuthHeatMapGrid_points):
            rangeAzimuthHeatMapGrid_xlin = np.append(rangeAzimuthHeatMapGrid_xlin, range_width)

        rangeAzimuthHeatMapGrid_ylin = np.arange(0, range_depth, 1.0 * range_depth / (rangeAzimuthHeatMapGrid_points - 1))
        if (len(rangeAzimuthHeatMapGrid_ylin) < rangeAzimuthHeatMapGrid_points):
            rangeAzimuthHeatMapGrid_ylin = np.append(rangeAzimuthHeatMapGrid_ylin, range_depth)


        QQ = np.array(QQ_)
        posX = np.array(posX_)
        posY = np.array(posY_)
        xlin = np.array(rangeAzimuthHeatMapGrid_xlin)
        ylin = np.array(rangeAzimuthHeatMapGrid_ylin)

        X, Y = np.meshgrid(xlin, ylin)

        FQQ = QQ[:, 1:]

        #左右反转图像数据
        FlipQQ = np.fliplr(FQQ)

        XY = np.hstack((posX.flatten()[:, None], posY.flatten()[:, None]))

        FQ = FlipQQ.flatten()

        Z = sci.griddata(XY, FQ, (X, Y), method="linear")

        #上下翻转图像数据
        Z = cv2.flip(Z, 0, dst=None)

        plt.axis('on')
        plt.imshow(Z, cmap=plt.cm.jet)
        plt.savefig("out.jpg")

        end = time.perf_counter()
        runTime = end - start
        print(runTime)

        img = cv2.imread("out.jpg")

        # 测试数据是否在变化
        # cv2.imshow("demo",img)
        # cv2.waitKey(200)

        img = cv2.copyMakeBorder(img,160,160,0,0,cv2.BORDER_REPLICATE)
        img_height, img_width = img.shape[0:2]

        # print(img_width)
        # print(img_height)

        width_crop = (img_width - 480) // 2

        if width_crop > 0:
            img = img[:, width_crop:-width_crop, :]

        # cv2.imshow("demo",img)
        # cv2.waitKey(0)

        yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV_I420)
        # cv2.imshow("demo",yuv)
        # key = cv2.waitKey(0)
        saveInYUV(yuv)
        tcpclient_SendData()


def saveInYUV(img_yuv):

    fp = open("out_in_yuv.yuv", "wb")

    yuv = np.array(img_yuv, "B")
    yuv.tofile(fp)
    fp.close()

def tcpserver_Getdata():
    
    server.listen(1)                     
    try:
        client,addr = server.accept()          
        print(addr,"Connected")

        data_slice = ""

        while True:
            data = client.recv(MaxBytes)

            if not data:
                # print(data_slice)
                print("data is Null")


                #返回一次帧的数据
                return data_slice

            localTime = time.asctime( time.localtime(time.time()))
            print(localTime,' Recv NumBytes:',len(data))

            # print(binascii.hexlify(data))

            string = str(binascii.hexlify(data))
            string = string[2:-1]

            data_slice = data_slice + string

            # client.send("OK!")
    except BaseException as e:
        print("Error: ")
        print(repr(e))
    finally:
        # server.close()                    
        print("\nquit\n")

def tcpclient_SendData():

    host_client_SendData = '192.168.3.69'
    port_client_SendData = 6666
    client = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    client.settimeout(30)
    client.connect((host_client_SendData, port_client_SendData))

    file_name = "out_in_yuv.yuv"

    fp = open(file_name, "rb")

    inputData = fp.read()

    try:
        sendBytes = client.sendall(inputData)
        
    except Exception as e:
        print(e)
    
    fp.close()
    client.close()


    print("\nquit\n")

if __name__ == "__main__":
     main()
